import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks


def get_model(n_cls, input_shape=(224, 224, 3)):
    inputs = keras.Input(shape=input_shape)
    base_model = keras.applications.MobileNet(
        input_shape=input_shape,
        alpha=1.0,
        depth_multiplier=1,
        dropout=1e-3,
        include_top=False,
        weights='imagenet',
        pooling='avg',
    )
    base_model.trainable = False
    x = base_model(inputs)
    customer_model = keras.layers.Dense(n_cls, activation='softmax')
    x = customer_model(x)
    model = keras.Model(inputs, x)
    return model


if '__main__' == __name__:
    import python_ai.common.read_data.cat_dog_teachers as cat_dog
    import os

    VER = 'v1.1'
    ALPHA = 1e-3
    SIZE = 224
    BATCH_SIZE = 32
    N_EPOCHS = 2
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    dir = '../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/train'
    IMG_DIR = BASE_DIR + '/' + dir
    LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)

    model = get_model(2, (SIZE, SIZE, 3))
    model.summary()

    (x_train, y_train), (x_val, y_val), (x_test, y_test) = cat_dog.load_data(IMG_DIR, 1, SIZE, SIZE, 'cats', 'dogs', ch_first=False)

    model.compile(
        loss=losses.sparse_categorical_crossentropy,
        optimizer=optimizers.Adam(learning_rate=ALPHA),
        metrics=[metrics.sparse_categorical_accuracy]
    )

    model.fit(x_train,
              y_train,
              BATCH_SIZE,
              N_EPOCHS,
              callbacks=[callbacks.TensorBoard(LOG_DIR, update_freq='batch', profile_batch=0)],
              validation_data=(x_val, y_val),
              validation_batch_size=BATCH_SIZE,
              )

    print('Testing...')
    model.evaluate(x_test,
                   y_test,
                   BATCH_SIZE,
                   )
    print('Tested.')
